from pymel.api import *
import maya.OpenMaya as om
import maya.cmds as cmd
from engine import *
import os
import sys
import math

# =================================================================================================
# Extract the meshdata
# =================================================================================================
def ExtractMeshData(dagpath, outputfile):
    # -- print the name and write
    print dagpath.partialPathName()
    outputfile.write(dagpath.partialPathName() + ":\n")

    # -- view node as mesh and transform
    mesh = om.MFnMesh(dagpath)
    trans = om.MFnTransform(dagpath)

    # -- get the transformation matrix
    mat = trans.transformation().asMatrix()

    # -- create c arrays
    vertices = om.MFloatPointArray()
    vertcount = om.MIntArray()
    indices = om.MIntArray()

    # -- get the properties
    mesh.getPoints(vertices)
    mesh.getVertices(vertcount, indices)

    # -- write the properties to a file
    title = "Vertices: %s" % vertices.length()
    print title
    outputfile.write(title + "\n")
    for i in range(vertices.length()):
        output = "%f %f %f " % (vertices[i].x, vertices[i].y, vertices[i].z)
        print output
        outputfile.write(output + "\n")

    # -- get the vertcount of the mesh
    title = "Vertcount: %s" % vertcount.length()
    print title
    outputfile.write(title + "\n")
    for i in range(vertcount.length()):
        output = "%d" % (vertcount[i]) 
        print output
        outputfile.write(output + "\n")

    # -- get the indices of the mesh
    title = "Indices: %s" % indices.length()
    print title
    outputfile.write(title + "\n") 
    for i in range(indices.length()):
        output = "%d" % (indices[i]) 
        print "%d" % (indices[i])
        outputfile.write(output + "\n")

    # -- get the transformation and transpose matrix since maya 
    # -- is in righthanded coordinate system and direct3d is not
    trans = 'WorldTransformation:\n' 
    for row in range(4):
        trans += '%g %g %g %g\n' % (mat(row,0), mat(row,1), mat(row,2), mat(row,3))

    outputfile.write(trans)

    # -- just for extra info print this too...
    current = 0
    for i in range(len(vertcount)):
        count = vertcount[i]
        print "from %d to %d " % (current, current + count)
        for j in range(current, current + count):
            print "%d, " % (indices[j])
            print "%f, %f, %f, " % (vertices[indices[j]].x, vertices[indices[j]].y, 
                                    vertices[indices[j]].z)
        current = current + count

    print trans

# =================================================================================================
# Extract the directional light data from the scene
# =================================================================================================
def ExtractDirLightData(dagpath, outputfile):
    # - print the name and write
    print dagpath.partialPathName()
    outputfile.write(dagpath.partialPathName() + ":\n")

    # -- get the light transformation
    trans = om.MFnTransform(dagpath)
    light = om.MFnLight(dagpath)

    # -- get the transformation matrix
    mat = trans.transformation().asMatrix()

    # -- get the intensity
    intense = 'Intensity:'
    print intense
    print light.intensity()

    outputfile.write(intense + '\n')
    outputfile.write(str(light.intensity()) + '\n')

    # -- get the color
    print 'Color:'
    outputfile.write('Color:\n')
    col = light.color()
    colvec = "%g %g %g %g" % (col.r, col.g, col.b, col.a)
    print colvec
    outputfile.write(colvec)

    # -- get the transformation and transpose matrix since maya 
    # -- is in righthanded coordinate system and direct3d is not
    trans = 'DirLightTransformation:\n' 
    for row in range(3):
        trans += '%g %g %g 0.0' % (mat(row,0), mat(row,1), mat(row,2))
        trans += '\n'

    # -- likewise negate x translation for the same reason as above
    trans += '%g %g %g 1.0\n'  % (0.0, 0.0, 0.0)

    outputfile.write(trans)

    print trans

# =================================================================================================
# Extract the point light data from the scene
# =================================================================================================
def ExtractPointLightData(dagpath, outputfile):
    # - print the name and write
    print dagpath.partialPathName()
    outputfile.write(dagpath.partialPathName() + ":\n")

    # -- get the light transformation
    trans = om.MFnTransform(dagpath)
    light = om.MFnLight(dagpath)

    # -- get the transformation matrix
    mat = trans.transformation().asMatrix()

    # -- get the intensity
    intense = 'Intensity:'
    print intense
    print light.intensity()
    outputfile.write(intense + '\n')
    outputfile.write(str(light.intensity()) + '\n')

    # -- get the color
    print 'Color:'
    outputfile.write('Color:\n')
    col = light.color()
    colvec = "%g %g %g %g" % (col.r, col.g, col.b, col.a)
    print colvec
    outputfile.write(colvec)

    # -- get the transformation and transpose matrix since maya 
    # -- is in righthanded coordinate system and direct3d is not
    trans = 'PointLightTransformation:\n' 
    for row in range(4):
        trans += '%g %g %g %g' % (mat(row,0), mat(row,1), mat(row,2), mat(row,3))
        trans += '\n'

    outputfile.write(trans)

    print trans

# =================================================================================================
# Export spot light data
# =================================================================================================
def ExtractSpotLightData(dagpath, outputfile):
    # - print the name and write
    print dagpath.partialPathName()
    outputfile.write(dagpath.partialPathName() + ":\n")

    # -- get the light transformation
    #$$$:NAB we'll worry about color later :P
    trans = om.MFnTransform(dagpath)
    light = om.MFnSpotLight(dagpath)

    # -- get the transformation matrix
    rotorder = om.MEulerRotation.kXYZ
    pos = trans.transformation().getTranslation(om.MSpace.kWorld)
    qrot = trans.transformation().rotation()
    rot = qrot.asEulerRotation()
    rot.reorderIt(0)

    # -- get the intensity
    intense = 'Intensity:'
    print intense
    print light.intensity()
    outputfile.write(intense + '\n')
    outputfile.write(str(light.intensity()) + '\n')

    # -- get the umbra
    umbra = 'Umbra: '
    umbraangle = math.cos(light.coneAngle()/2.0)
    print umbra
    print umbraangle
    outputfile.write(str(umbra) + '\n')
    outputfile.write(str(umbraangle) + '\n')

    # -- get the penumbra
    penum = 'Penumbra: '
    penumangle = math.cos(light.penumbraAngle()/2.0)
    print penum
    print penumangle
    outputfile.write(str(penum) + '\n')
    outputfile.write(str(penumangle) + '\n')

    # -- get the color
    print 'Color:'
    outputfile.write('Color:\n')
    col = light.color()
    colvec = "%g %g %g %g" % (col.r, col.g, col.b, col.a)
    print colvec
    outputfile.write(str(colvec) + '\n')

    # -- get the transformation and transpose matrix since maya 
    # -- is in righthanded coordinate system and direct3d is not
    trans = 'SpotLightOrientation:' 
    rowvec = "%f %f %f" % (rot.x, -rot.y, rot.z) 
    posrow = "%f %f %f" % (pos.x, pos.y, pos.z)
    print trans
    print rowvec
    print posrow
    outputfile.write(str(trans) + '\n')
    outputfile.write(str(rowvec) + '\n')
    outputfile.write(str(posrow) + '\n')

# =================================================================================================
# Export the specific type and write as the name
# =================================================================================================
def ExportType(outputfile, iterator, name, *types):
    # -- while there is still more to search
    # -- look for node data

    i = 0
    while not iterator.isDone():
        hasallatr = True
        dagpath = om.MDagPath()
        iterator.getPath(dagpath)
        for attrs in types:
            hasallatr &= dagpath.hasFn(attrs)
        if hasallatr:
            # -- begin node exporting
            outputfile.write("%s:\n" % name)
            getattr(sys.modules[__name__],"Extract%sData" % name)(dagpath, outputfile)
            iterator.next()
            i += 1
        else:
            iterator.next()
            continue

    iterator.reset()

# =================================================================================================
# Export resources from a maya binary
# =================================================================================================
def Export(filepath):
    filepath.strip(".mb")
    
    fileloc = path.GetMayaFileLoc()

    # -- open maya file
    MFileIO.open(fileloc + "/" + filepath + ".mb")

    root = path.GetRoot()
    meshdir = os.path.join(root, "resources/objects/")
    existence = os.path.isdir(meshdir)

    if not existence:
        print "%s does not exist!" % meshdir
        print "Exiting script."
        return -1
    else:
        outputfile = file(meshdir + "object.txt", "w")

    # -- create the iterator
    iterator = om.MItDag(om.MItDag.kBreadthFirst, om.MFn.kInvalid)

    # -- export the mesh data
    ExportType(outputfile, iterator, "Mesh", om.MFn.kMesh, om.MFn.kTransform)

    # -- export the direction light data
    ExportType(outputfile, iterator, "DirLight", om.MFn.kDirectionalLight, om.MFn.kTransform)

    # -- export the point light data
    ExportType(outputfile, iterator, "PointLight", om.MFn.kPointLight, om.MFn.kTransform)

    # -- export the spot light data
    ExportType(outputfile, iterator, "SpotLight", om.MFn.kSpotLight, om.MFn.kTransform)


    cmd.quit()
    
